package com.example.springai.repository;

import com.example.springai.bean.ChatMessage;
import com.example.springai.mapper.ChatMessageMapper;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

@Service
public class InDatabaseChatMemory implements ChatMemory {

    @Autowired
    private ChatMessageMapper chatMessageMapper;

    @Override
    public void add(String conversationId, List<Message> messages) {
        for (Message message : messages) {
            String role = "";
            switch (message.getMessageType()){
                case USER -> role = "user";
                case ASSISTANT ->  role = "assistant";
                default -> role = "unknown";
            }

            ChatMessage chatMessage = new ChatMessage();
            chatMessage.setConversationId(conversationId);
            chatMessage.setRole(role);
            chatMessage.setContent(message.getText());

            // 保存到数据库
            chatMessageMapper.save(chatMessage);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {

        List<ChatMessage> chatMessages = chatMessageMapper.findByConversationId(conversationId);
        List<Message> messages = new ArrayList<>();
        for (ChatMessage chatMessage : chatMessages) {
            switch (chatMessage.getRole()){
                case "user" -> messages.add(new UserMessage(chatMessage.getContent()));
                case "assistant" -> messages.add(new AssistantMessage(chatMessage.getContent()));
                default -> throw new RuntimeException("未知的角色");
            }
        }
        return messages;
    }

    @Override
    public void clear(String conversationId) {
        chatMessageMapper.deleteByConversationId(conversationId);
    }
}
